package cim144.ctdp
import Chisel._
import chisel3.{Mux, RegInit, RegNext, SInt, UInt, Vec, VecInit, Wire, when}
import chisel3.util.{Cat, RegEnable, log2Ceil}
import cim144.{SystemConfig, cim_mvm_8bit_shift, mvm_IO}

object mvmcim {
  def apply(custom_mvm:Bool,ismvm:Bool,bitwise_end:Bool,
            adc_range:UInt,out_shift:UInt,push_1152:UInt,
            r1:Bits,r2:Bits
           ):(Bool,Bool,UInt)={
    val inst = Module(new mvmcim()).io
    inst.custom_mvm   :=custom_mvm
    inst.ismvm        :=ismvm
    inst.bitwise_end  :=bitwise_end
    inst.adc_range    :=adc_range
    inst.out_shift    :=out_shift
    inst.push_1152    :=push_1152
    inst.r1           :=r1
    inst.r2           :=r2
    (inst.cim_done,inst.mvm_done,inst.save_1024)
  }
}


class mvmcim extends Module with SystemConfig{
  val io = new Bundle{
    val custom_mvm  = Input(Bool())
    val ismvm       = Input(Bool())
    val bitwise_end = Input(Bool())
    val adc_range   = Input(UInt(8.W))
    val out_shift   = Input(UInt(8.W))
    val push_1152   = Input(UInt(1152.W))
    val r1          = Bits(INPUT,64)
    val r2          = Bits(INPUT,64)

    val cim_done    = Output(Bool())
    val mvm_done    = Output(Bool())
    val save_1024   = Output(UInt((SAVE_MAX_SIZE*CIM_XLEN).W))
  }
  val mvm_start_tmp = RegNext(Mux(io.custom_mvm,true.B,Mux(io.cim_done&io.ismvm,true.B,false.B)))
  val mvm_start     = RegNext(Mux(io.mvm_done,false.B,mvm_start_tmp))
  val rcbd = Wire(new mvm_IO())
  //rs1 |31******************16|15***************0|
  //    |*******row_num********|******col_num*****|
  //rs2 |31*********22|21*********11|10**********0|
  //    |*array index*|**start_row**|**start_col**|
  rcbd.row_begin := RegEnable(io.r2(21,11),0.U,io.custom_mvm)
  rcbd.row_end   := RegEnable(io.r1(31,16)+io.r2(21,11) - 1.U,0.U,io.custom_mvm)
  rcbd.col_begin := RegEnable(io.r2(10, 0),0.U,io.custom_mvm)
  rcbd.col_end   := RegEnable(io.r1(15, 0)+io.r2(10, 0) - 1.U,0.U,io.custom_mvm)
  // *********** 1. get cim data
  val cim_out     = cim_mvm_8bit_shift(mvm_start,rcbd,io.push_1152,io.adc_range)
  io.cim_done    := RegNext(cim_out._1)
  val cim_result = cim_out._2.asTypeOf(Vec(COL_NUM,SInt(8.W)))
  // ************ 2. Acc
  val mvm_tmp         = RegInit(VecInit(Seq.fill(COL_NUM)(0.S(16.W))))
  val mvm_tmp_shift   = VecInit(Seq.fill(COL_NUM)(0.S(16.W)))
  val mvm_result      = VecInit(Seq.fill(COL_NUM)(0.S(8.W)))
  for(i<-0 until COL_NUM) {
    when(io.custom_mvm) {
      mvm_tmp(i) := 0.S
    }.elsewhen(io.cim_done) {
      mvm_tmp(i) := mvm_tmp(i) + cim_result(i)
    }
    mvm_tmp_shift(i) := mvm_tmp(i) >> io.out_shift
    mvm_result(i) := Mux(mvm_tmp_shift(i) >= 127.S, 127.S, Mux(mvm_tmp_shift(i) <= -127.S, -127.S, mvm_tmp_shift(i)))
  }
  io.mvm_done := io.ismvm & io.bitwise_end
  //io.save_1024 = mvm_result.asTypeOf(Vec(SAVE_MAX_SIZE,UInt(CIM_XLEN.W)))
  io.save_1024 := mvm_result.asUInt()
}
